#include "Skeleton.h"
#include <cmath>
#include <algorithm>
#include <unordered_map>
#include <fstream>
#include <iomanip> // for std::setprecision

namespace {
    // 并查集查找（带路径压缩）
    int find(std::vector<int>& parent, int i) {
        if (parent[i] != i)
            parent[i] = find(parent, parent[i]);
        return parent[i];
    }

    // 并查集合并
    void unite(std::vector<int>& parent, int i, int j) {
        int ri = find(parent, i);
        int rj = find(parent, j);
        if (ri != rj)
            parent[ri] = rj;
    }

    // 欧几里得距离
    double distance(const Point& a, const Point& b) {
        double dx = a.first - b.first;
        double dy = a.second - b.second;
        return std::sqrt(dx * dx + dy * dy);
    }
}

std::vector<EdgeSegment> computeMST(const std::vector<Point>& points) {
    const int n = points.size();
    std::vector<int> parent(n);
    for (int i = 0; i < n; ++i)
        parent[i] = i;

    // 1. 构造所有边及其长度
    struct Edge {
        int from;
        int to;
        double weight;
    };
    std::vector<Edge> edges;
    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            edges.push_back({i, j, distance(points[i], points[j])});
        }
    }

    // 2. 边按权重排序
    std::sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b) {
        return a.weight < b.weight;
    });

    // 3. Kruskal 主体：选边且不形成环
    std::vector<EdgeSegment> mst;
    for (const auto& edge : edges) {
        if (find(parent, edge.from) != find(parent, edge.to)) {
            unite(parent, edge.from, edge.to);
            mst.emplace_back(points[edge.from], points[edge.to]);
        }
    }

    return mst;
}


// 骨架线段自适应加密
std::vector<EdgeSegment> densifySkeletonEdgesAdaptive(const std::vector<EdgeSegment>& edges, double maxSegmentLength) {
    std::vector<EdgeSegment> densified;
    for (const auto& [p1, p2] : edges) {
        double len = distance(p1, p2);
        int segments = std::ceil(len / maxSegmentLength);
        for (int i = 0; i < segments; ++i) {
            double t1 = static_cast<double>(i) / segments;
            double t2 = static_cast<double>(i + 1) / segments;
            Point q1 = {p1.first * (1 - t1) + p2.first * t1, p1.second * (1 - t1) + p2.second * t1};
            Point q2 = {p1.first * (1 - t2) + p2.first * t2, p1.second * (1 - t2) + p2.second * t2};
            densified.emplace_back(q1, q2);
        }
    }
    return densified;
}


// 提取骨架线段
std::unordered_map<int, std::vector<EdgeSegment>> extractSkeleton(
    const std::vector<std::vector<Edge>>& clusteredEdges,
    const std::vector<Node>& nodes,
    double maxSegmentLength 
) {
    // 创建节点查找表
    std::unordered_map<int, Node> nodeMap;
    for (const auto& node : nodes) {
        nodeMap[node.id] = node;
    }

    std::unordered_map<int, std::vector<EdgeSegment>> skeletons;

    for (int clusterId = 0; clusterId < clusteredEdges.size(); ++clusterId) {
        const auto& edges = clusteredEdges[clusterId];

        std::vector<Point> midpoints;
        for (const auto& edge : edges) {
            const auto& p1 = nodeMap[edge.from];
            const auto& p2 = nodeMap[edge.to];
            midpoints.emplace_back(
                (p1.lon + p2.lon) / 2.0,
                (p1.lat + p2.lat) / 2.0
            );
        }

        auto mst = computeMST(midpoints);
        auto densified = densifySkeletonEdgesAdaptive(mst, maxSegmentLength);
        skeletons[clusterId] = densified;
    }

    return skeletons;
}


void writeSkeletonToCsv(const std::unordered_map<int, std::vector<std::pair<Point, Point>>>& skeletons, const std::string& filepath) {
    std::ofstream out(filepath);
    if (!out) {
        throw std::runtime_error("Failed to open file for writing: " + filepath);
    }

    // 写入 header
    out << "cluster_id,x1,y1,x2,y2\n";
    out << std::fixed << std::setprecision(8); // 控制浮点数格式

    for (const auto& [clusterId, segments] : skeletons) {
        for (const auto& [p1, p2] : segments) {
            out << clusterId << ","
                << p1.first << "," << p1.second << ","
                << p2.first << "," << p2.second << "\n";
        }
    }

    out.close();
}